#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import yaml
import re
import sys
import os
import json
import subprocess
from pathlib import Path

def get_ansible_inventory(inventory_file):
    """
    使用 ansible-inventory 命令获取解析后的库存信息
    返回包含所有主机和变量信息的JSON字典
    """
    try:
        # 执行 ansible-inventory 命令
        cmd = ['ansible-inventory', '-i', inventory_file, '--list']
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        # 解析JSON输出
        inventory_data = json.loads(result.stdout)
        return inventory_data
    except subprocess.CalledProcessError as e:
        print(f"执行 ansible-inventory 命令时出错: {e}")
        print(f"错误输出: {e.stderr}")
        return {}
    except json.JSONDecodeError as e:
        print(f"解析JSON输出时出错: {e}")
        return {}
    except Exception as e:
        print(f"获取Ansible库存信息时发生未知错误: {e}")
        return {}

def update_mode_based_on_host_count(inventory_data):
    """
    根据主机数量更新库存数据中的模式变量。
    如果组件有多个主机节点，设置模式为cluster；否则为single。
    """
    # 获取所有组，排除特殊键
    special_keys = ['all', 'ungrouped', '_meta']
    groups = {}
    for key, value in inventory_data.items():
        if key not in special_keys and not key.startswith('_'):
            groups[key] = value

    hostvars = inventory_data.get('_meta', {}).get('hostvars', {})

    for group_name, group_data in groups.items():
        if 'hosts' in group_data:
            host_count = len(group_data['hosts'])
            mode_value = 'cluster' if host_count > 1 else 'single'
            mode_var_name = f"{group_name}_mode"

            # 更新每个主机的这个变量
            for host in group_data['hosts']:
                if host in hostvars:
                    hostvars[host][mode_var_name] = mode_value
                else:
                    print(f"警告: 主机 {host} 的变量未找到在hostvars中")

    return inventory_data

def replace_jinja_variables(data, inventory_data, present_components):
    """
    递归遍历YAML数据，替换所有Jinja2变量引用（{{ ... }} 和 hostvars）
    使用从ansible-inventory获取的数据
    新增参数 present_components: 存在于inventory中的组件列表
    """
    if isinstance(data, dict):
        for key, value in data.items():
            data[key] = replace_jinja_variables(value, inventory_data, present_components)
        return data
    elif isinstance(data, list):
        return [replace_jinja_variables(item, inventory_data, present_components) for item in data]
    elif isinstance(data, str):
        # 替换 hostvars['some_host'].some_var - 保持原逻辑，因为hostvars已指定主机
        hostvars_pattern = r"\{\{\s*hostvars\['(\w+)'\]\.(\w+)\s*\}\}"
        def replace_hostvars(match):
            host_name = match.group(1)
            var_name = match.group(2)

            # 从inventory数据中查找主机变量
            host_vars = inventory_data.get('_meta', {}).get('hostvars', {}).get(host_name, {})
            if var_name in host_vars:
                return str(host_vars[var_name])
            else:
                print(f"警告: 无法找到 hostvars['{host_name}'].{var_name}，保留原值")
                return match.group(0)

        data = re.sub(hostvars_pattern, replace_hostvars, data)

        # 替换简单的 {{ some_var }}
        simple_var_pattern = r"\{\{\s*(\w+)\s*\}\}"
        def replace_simple_var(match):
            var_name = match.group(1)

            # 新增逻辑: 检查变量是否属于存在的组件
            # 假设变量命名约定为 {组件名}_{变量名} (例如 mysql_mode)
            var_belongs_to_present_component = False
            for comp in present_components:
                if var_name.startswith(comp + "_"):
                    var_belongs_to_present_component = True
                    break

            if not var_belongs_to_present_component:
                # 变量不属于任何存在的组件，保留原值
                print(f"信息: 变量 {var_name} 不属于任何存在的组件，保留原值")
                return match.group(0)

            # 变量属于存在的组件，尝试在hostvars中查找
            for host, host_vars in inventory_data.get('_meta', {}).get('hostvars', {}).items():
                if var_name in host_vars:
                    return str(host_vars[var_name])

            print(f"警告: 简单变量 {var_name} 未在hostvars中找到，保留原值")
            return match.group(0)

        data = re.sub(simple_var_pattern, replace_simple_var, data)

        return data
    else:
        return data

def resolve_simple_variables(data):
    """
    解析简单的变量引用（{{ some_var }}），这些变量通常在YAML文件自身中定义。
    这是一个简单的实现，假设变量定义在顶层且没有循环引用。
    """
    if isinstance(data, dict):
        # 首先收集所有顶层的变量定义
        variables = {}
        for key, value in data.items():
            if not isinstance(value, dict) and not isinstance(value, list):
                variables[key] = value

        # 然后递归替换
        for key, value in data.items():
            if isinstance(value, dict):
                data[key] = resolve_simple_variables(value)
            elif isinstance(value, list):
                data[key] = [resolve_simple_variables(item) for item in value]
            elif isinstance(value, str):
                # 替换 {{ varname }}
                var_pattern = r"\{\{\s*(\w+)\s*\}\}"
                while re.search(var_pattern, value):
                    match = re.search(var_pattern, value)
                    var_name = match.group(1)
                    if var_name in variables:
                        value = value.replace(match.group(0), str(variables[var_name]))
                    else:
                        print(f"警告: 未找到变量 {var_name} 的定义，保留原值")
                        break
                data[key] = value
        return data
    elif isinstance(data, list):
        return [resolve_simple_variables(item) for item in data]
    else:
        return data

def main():
    # 检查命令行参数数量
    if len(sys.argv) < 4:
        print("Usage: python script.py <root_path> <inventory_file> <custom_file>")
        print("Example: python script.py /path/to/your/dir inventory.ini custom.yml")
        sys.exit(1)

    # 文件路径
    root_path = sys.argv[1]
    inventory_file = sys.argv[2]
    custom_file = sys.argv[3]
    output_file = os.path.join(root_path, "custom_resolved.yml")

    print(f"正在使用ansible-inventory解析库存文件: {inventory_file}")
    inventory_data = get_ansible_inventory(inventory_file)

    if not inventory_data:
        print("未能从ansible-inventory获取库存信息，请检查文件格式和ansible安装。")
        return

    print("成功获取Ansible库存信息")

    # 根据主机数量更新模式变量
    print("根据主机数量更新模式变量...")
    inventory_data = update_mode_based_on_host_count(inventory_data)

    # 新增: 获取存在的组件列表 (非特殊组且包含主机)
    present_components = []
    for group_name, group_data in inventory_data.items():
        if group_name not in ['all', 'ungrouped', '_meta'] and isinstance(group_data, dict) and 'hosts' in group_data and len(group_data['hosts']) > 0:
            present_components.append(group_name)
    print(f"存在的组件: {present_components}")

    print(f"正在读取custom.yml文件: {custom_file}")
    try:
        with open(custom_file, 'r', encoding='utf-8') as f:
            custom_data = yaml.safe_load(f)
    except Exception as e:
        print(f"读取或解析YAML文件时出错: {e}")
        return

    print("正在替换变量引用...")
    # 使用ansible-inventory数据替换变量引用，并传入存在的组件列表
    custom_data = replace_jinja_variables(custom_data, inventory_data, present_components)
    # 然后替换简单变量引用 (YAML内部定义的变量)
    custom_data = resolve_simple_variables(custom_data)

    print(f"正在写入固化后的YAML文件: {output_file}")
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            yaml.dump(custom_data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
        print(f"成功! 固化后的配置文件已保存至: {output_file}")
    except Exception as e:
        print(f"写入文件时出错: {e}")

if __name__ == "__main__":
    main()
